import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import time
from typing import List, Dict, Tuple

class BattleAnimator:
    """战斗动画可视化工具"""
    
    def __init__(self, width: int = 800, height: int = 600):
        """初始化战斗动画器"""
        self.width = width
        self.height = height
        self.fig, self.ax = plt.subplots(figsize=(width/100, height/100))
        self.ax.set_xlim(0, width)
        self.ax.set_ylim(0, height)
        self.ax.axis("off")
        
        self.allies = []  # 我方单位
        self.enemies = []  # 敌方单位
        self.effects = []  # 特效
        self.anim = None
        self.running = False
    
    def add_ally(self, unit_id: str, x: float, y: float, name: str, health: int = 100):
        """添加我方单位"""
        unit = {
            "id": unit_id,
            "x": x,
            "y": y,
            "name": name,
            "health": health,
            "max_health": health,
            "sprite": self.ax.scatter([x], [y], c="blue", s=100, label=name)
        }
        self.allies.append(unit)
        self.ax.text(x, y-30, name, ha="center", va="center", color="blue")
        return unit
    
    def add_enemy(self, unit_id: str, x: float, y: float, name: str, health: int = 100):
        """添加敌方单位"""
        unit = {
            "id": unit_id,
            "x": x,
            "y": y,
            "name": name,
            "health": health,
            "max_health": health,
            "sprite": self.ax.scatter([x], [y], c="red", s=100, label=name)
        }
        self.enemies.append(unit)
        self.ax.text(x, y-30, name, ha="center", va="center", color="red")
        return unit
    
    def _update_health(self, unit):
        """更新单位生命值显示"""
        # 清除旧的生命值条
        for artist in self.ax.texts:
            if artist.get_text().startswith(f"{unit['name']}_health"):
                artist.remove()
        
        # 绘制新的生命值条
        health_percent = unit["health"] / unit["max_health"]
        health_color = "green" if health_percent > 0.5 else "yellow" if health_percent > 0.2 else "red"
        
        # 生命值文本
        self.ax.text(
            unit["x"], unit["y"] + 30,
            f"{unit['health']}/{unit['max_health']}",
            ha="center", va="center", 
            bbox=dict(facecolor=health_color, alpha=0.5),
            name=f"{unit['name']}_health"
        )
    
    def attack(self, attacker_id: str, target_id: str, damage: int) -> bool:
        """执行攻击动画"""
        # 查找攻击者和目标
        attacker = None
        for unit in self.allies + self.enemies:
            if unit["id"] == attacker_id:
                attacker = unit
                break
        
        target = None
        for unit in self.allies + self.enemies:
            if unit["id"] == target_id:
                target = unit
                break
        
        if not attacker or not target:
            return False
        
        # 添加攻击特效
        effect = self.ax.plot(
            [attacker["x"], target["x"]], 
            [attacker["y"], target["y"]], 
            "orange", linestyle="--", linewidth=2
        )
        self.effects.append({"effect": effect, "frame": 0})
        
        # 目标受伤
        target["health"] = max(0, target["health"] - damage)
        self._update_health(target)
        
        return True
    
    def _animate(self, frame):
        """动画更新函数"""
        # 更新特效（短暂显示后消失）
        to_remove = []
        for i, effect in enumerate(self.effects):
            effect["frame"] += 1
            if effect["frame"] > 5:  # 显示5帧后移除
                for line in effect["effect"]:
                    line.remove()
                to_remove.append(i)
        
        # 移除过期特效
        for i in reversed(to_remove):
            self.effects.pop(i)
        
        return [sprite for unit in self.allies + self.enemies for sprite in unit["sprite"]]
    
    def start_battle(self):
        """开始战斗动画"""
        self.running = True
        self.anim = animation.FuncAnimation(
            self.fig, self._animate, interval=100, blit=True
        )
        plt.tight_layout()
        plt.show()
    
    def end_battle(self):
        """结束战斗动画"""
        self.running = False
        if self.anim:
            self.anim.event_source.stop()
        plt.close(self.fig)

# 使用示例
if __name__ == "__main__":
    animator = BattleAnimator()
    
    # 添加单位
    animator.add_ally("a1", 200, 300, "风堇", 120)
    animator.add_ally("a2", 200, 400, "缇宝", 100)
    animator.add_enemy("e1", 600, 300, "来古士", 150)
    animator.add_enemy("e2", 600, 400, "杂兵", 80)
    
    # 启动动画（在单独线程中）
    import threading
    threading.Thread(target=animator.start_battle, daemon=True).start()
    
    # 模拟战斗流程
    time.sleep(1)
    animator.attack("a1", "e2", 30)
    time.sleep(1)
    animator.attack("e1", "a2", 25)
    time.sleep(1)
    animator.attack("a2", "e1", 40)
    time.sleep(2)
    animator.end_battle()